#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *

parser = argparse.ArgumentParser(description='Plot test result')
parser.add_argument('-i', '--input', help='Input data file', default='hs.360x180x26.h0.nc')
parser.add_argument('-o', '--output', help='Output figure')
args = parser.parse_args()

fig = plt.figure(figsize=(12, 8))

f = addfile(args.input).isel(time=slice(200, -1))

ax = fig.add_subplot(2, 3, 1)
if not os.path.isfile('u_zavg.nc'):
	u_zavg = f.u.mean(dim=('time', 'lon'))
	u_zavg.to_netcdf('u_zavg.nc')
else:
	u_zavg = xr.open_dataarray('u_zavg.nc')
u_zavg.attrs['long_name'] = 'Zonal wind (m s^-1)'
levels = np.linspace(-15, 35, 11)
colors = ListedColormap([
	'#2A64F6',
	'#76FBFD',
	'#B2FDFE',
	'#7CFB4C',
	'#D7FE51',
	'#FFFF54',
	'#FBE94E',
	'#EF8733',
	'#EB4626',
	'#EA3323'
])
plot_contour_lat(ax, u_zavg, levels=levels, cmap=colors, invert_yaxis=True, with_contour=True)

ax = fig.add_subplot(2, 3, 2)
if not os.path.isfile('t_zavg.nc'):
	t_zavg = f.t.mean(dim=('time', 'lon'))
	t_zavg.to_netcdf('t_zavg.nc')
else:
	t_zavg = xr.open_dataarray('t_zavg.nc')
t_zavg.attrs['long_name'] = 'Temperature (K)'
levels = np.linspace(180, 310, 14)
colors = ListedColormap([
	'#2A64F6',
	'#76FBFD',
	'#92FCFD',
	'#92FCA4',
	'#7CFB4C',
	'#A0FC4E',
	'#E4F650',
	'#FFFF54',
	'#FBE94E',
	'#F19F39',
	'#ED702D',
	'#EB4626',
	'#EA3323'
])
plot_contour_lat(ax, t_zavg, levels=levels, cmap=colors, invert_yaxis=True, with_contour=True)

u_mean = f.u.mean(dim='time')
v_mean = f.v.mean(dim='time')
t_mean = f.t.mean(dim='time')

ax = fig.add_subplot(2, 3, 3)
if not os.path.isfile('uv_adv_pert.nc'):
	uv_adv_pert = ((f.u - u_mean) * (f.v - v_mean)).mean(dim=('time', 'lon'))
	uv_adv_pert.to_netcdf('uv_adv_pert.nc')
else:
	uv_adv_pert = xr.open_dataarray('uv_adv_pert.nc')
uv_adv_pert.attrs['long_name'] = 'Eddy momentum flux (m^2 s^-2)'
levels = np.linspace(-80, 80, 17)
colors = ListedColormap([
	'#2A64F6',
	'#61D2FA',
	'#92FCFD',
	'#B2FDFE',
	'#92FCA4',
	'#7CFB4C',
	'#A0FC4E',
	'#D7FE51',
	'#FFFF54',
	'#FFFF54',
	'#FBE94E',
	'#F19F39',
	'#EF8733',
	'#ED702D',
	'#EA3F25',
	'#EA3323'
])
plot_contour_lat(ax, uv_adv_pert, levels=levels, cmap=colors, invert_yaxis=True, with_contour=True)

ax = fig.add_subplot(2, 3, 4)
if not os.path.isfile('ke_pert.nc'):
	ke_pert = ((f.u - u_mean)**2 + (f.v - v_mean)**2).mean(dim=('time', 'lon')) / 2
	ke_pert.to_netcdf('ke_pert.nc')
else:
	ke_pert = xr.open_dataarray('ke_pert.nc')
ke_pert.attrs['long_name'] = 'Eddy kinetic energy (m^2 s^-2)'
levels = np.linspace(-30, 420, 16)
colors = ListedColormap([
	'#2A64F6',
	'#61D2FA',
	'#92FCFD',
	'#B2FDFE',
	'#77FB54',
	'#7CFB4C',
	'#D7FE51',
	'#E4F650',
	'#FFFF54',
	'#FBE94E',
	'#F7CF46',
	'#EF8733',
	'#ED702D',
	'#EA3F25',
	'#EA3323'
])
plot_contour_lat(ax, ke_pert, levels=levels, cmap=colors, invert_yaxis=True, with_contour=True)

ax = fig.add_subplot(2, 3, 5)
if not os.path.isfile('vt_pert.nc'):
	vt_pert = ((f.t - t_mean) * (f.v - v_mean)).mean(dim=('time', 'lon'))
	vt_pert.to_netcdf('vt_pert.nc')
else:
	vt_pert = xr.open_dataarray('vt_pert.nc')
vt_pert.attrs['long_name'] = 'Eddy heat flux (K m s^-1)'
levels = np.linspace(-23, 22, 16)
colors = ListedColormap([
	'#2A64F6',
	'#61D2FA',
	'#92FCFD',
	'#B2FDFE',
	'#77FB54',
	'#7CFB4C',
	'#D7FE51',
	'#E4F650',
	'#FFFF54',
	'#FBE94E',
	'#F7CF46',
	'#EF8733',
	'#ED702D',
	'#EA3F25',
	'#EA3323'
])
plot_contour_lat(ax, vt_pert, levels=levels, cmap=colors, invert_yaxis=True, with_contour=True)

ax = fig.add_subplot(2, 3, 6)
if not os.path.isfile('tvar_pert.nc'):
	tvar_pert = ((f.t - t_mean)**2).mean(dim=('time', 'lon'))
	tvar_pert.to_netcdf('tvar_pert.nc')
else:
	tvar_pert = xr.open_dataarray('tvar_pert.nc')
tvar_pert.attrs['long_name'] = 'Temperature variance (K^2)'
levels = np.linspace(-4, 48, 14)
colors = ListedColormap([
	'#2A64F6',
	'#76FBFD',
	'#92FCFD',
	'#92FCA4',
	'#7CFB4C',
	'#A0FC4E',
	'#E4F650',
	'#FFFF54',
	'#FBE94E',
	'#F19F39',
	'#ED702D',
	'#EB4626',
	'#EA3323'
])
plot_contour_lat(ax, tvar_pert, levels=levels, cmap=colors, invert_yaxis=True, with_contour=True)

plt.tight_layout()
plt.show()
